import os
import csv
from typing import Iterable
from neo4j import GraphDatabase, Driver
from dotenv import load_dotenv


def get_driver() -> Driver:
    load_dotenv()
    uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
    user = os.getenv("NEO4J_USER", "neo4j")
    password = os.getenv("NEO4J_PASSWORD", "neo4j")
    return GraphDatabase.driver(uri, auth=(user, password))


SCHEMA_QUERIES = [
    # Use composite indexes for faster MERGE
    "CREATE CONSTRAINT IF NOT EXISTS FOR (f:Function) REQUIRE f.id IS UNIQUE",
]


def ensure_schema(driver: Driver) -> None:
    with driver.session() as session:
        for q in SCHEMA_QUERIES:
            session.execute_write(lambda tx: tx.run(q))


def rows_from_csv(csv_path: str) -> Iterable[dict]:
    with open(csv_path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            yield row


def make_function_id(qualified_name: str) -> str:
    # Qualified name is already close to a stable identifier
    return qualified_name


def import_callgraph_csv(driver: Driver, csv_path: str, batch_size: int = 5000) -> None:
    def parse_row(row: dict):
        caller_q = row.get("caller_qualified", "")
        callee_q = row.get("callee_qualified", "")
        return {
            "caller_id": make_function_id(caller_q),
            "caller_name": caller_q.split("::")[-1] if caller_q else "",
            "caller_file": row.get("caller_file", ""),
            "callee_id": make_function_id(callee_q),
            "callee_name": callee_q.split("::")[-1] if callee_q else "",
            "callee_file": row.get("callee_file", ""),
            "call_file": row.get("call_file", ""),
            "call_line": int(row.get("call_line", 0) or 0),
        }

    buffer = []
    with driver.session() as session:
        ensure_schema(driver)

        def flush(buf):
            if not buf:
                return
            session.execute_write(_ingest_batch, buf)
            buf.clear()

        for row in rows_from_csv(csv_path):
            buffer.append(parse_row(row))
            if len(buffer) >= batch_size:
                flush(buffer)
        flush(buffer)


def _ingest_batch(tx, rows: list[dict]):
    query = """
    UNWIND $rows AS r
    MERGE (caller:Function {id: r.caller_id})
      ON CREATE SET caller.name = r.caller_name, caller.file = r.caller_file
      ON MATCH  SET caller.name = coalesce(caller.name, r.caller_name), caller.file = coalesce(caller.file, r.caller_file)
    MERGE (callee:Function {id: r.callee_id})
      ON CREATE SET callee.name = r.callee_name, callee.file = r.callee_file
      ON MATCH  SET callee.name = coalesce(callee.name, r.callee_name), callee.file = coalesce(callee.file, r.callee_file)
    MERGE (caller)-[e:CALLS]->(callee)
      ON CREATE SET e.file = r.call_file, e.line = r.call_line
      ON MATCH  SET e.file = coalesce(e.file, r.call_file), e.line = coalesce(e.line, r.call_line)
    """
    tx.run(query, rows=rows)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="将调用关系 CSV 导入 Neo4j")
    parser.add_argument("--csv", required=True, help="CodeQL 查询生成的 CSV 文件")
    args = parser.parse_args()

    drv = get_driver()
    try:
        import_callgraph_csv(drv, args.csv)
    finally:
        drv.close()

